from typing import Any

from langchain_core.prompts import PromptTemplate

from app.decorator.timeit import timeit

rag_question_answer_prompt_template = """
# 角色设定
你是一个专业问答助手，优先根据提供的上下文信息回答问题。

# 任务要求
1. 如果上下文信息中已包含答案，直接回答。
2. 只在必要时调用 MCP 工具，不要编造答案。
3. 如果通过上下文或 MCP 工具仍无法获取答案，请回复："暂时还无法回答该问题"。

# 上下文信息：
{context}
"""

prompt_template_by_mcp_result = """
你是一个专业问答助手，专门通过获取结构化数据回答问题。请严格遵守以下规则：

1. 使用上下文中提供的结构化数据，回答问题
2. 精确解析JSON数据结构，对时间/金额等特殊字段自动格式化, 保持数值计算的准确性
3. 禁止虚构数据，若数据不足，必须明确告知无法回答

# 结构化数据：
{context}
"""


class QuestionAnswerPrompt:
    """
    智能问答Prompt
    """

    @timeit
    def __init__(self, query: str, context: list[str],
                 prompt_template: str = rag_question_answer_prompt_template) -> None:
        """
        QuestionAnswerPrompt 初始化
        :param query: 用户输入query
        :param context: 上下文信息
        """
        self.query = query
        self.context = context
        self.prompt_tem = prompt_template
        self.prompt = self._format_prompt()

    def _format_prompt(self) -> str:
        """
        格式化prompt
        :return: str
        """
        prompt_template = PromptTemplate.from_template(rag_question_answer_prompt_template)
        context = "\n".join(text for i, text in enumerate(self.context))
        return prompt_template.format(question=self.query, context=context)


if __name__ == '__main__':
    prompt = QuestionAnswerPrompt(query="使用Java语言进行性能优化", context=[
        "Java程序性能优化，让你的Java程序更快，更稳定",
        "学习了Java程序，为什么不学习Python"])

    print(f"prompt: {prompt.prompt}")

    prompt_template = PromptTemplate.from_template(rag_question_answer_prompt_template)
    result = prompt_template.format(question="使用Java语言进行性能优化", context=[
        "Java程序性能优化，让你的Java程序更快，更稳定",
        "学习了Java程序，为什么不学习Python"])
    print(f"result: {result}")
